﻿#include "dgsopengl.h"
#include<qdatastream.h>
#include <cstddef>
//#include<auxiliary.h>
# define CUDA_SAFE_CALL_ALWAYS(A) \
A; \
cudaDeviceSynchronize(); \
if (cudaPeekAtLastError() != cudaSuccess) \
qDebug() << cudaGetErrorString(cudaGetLastError());

#if DEBUG || _DEBUG
# define CUDA_SAFE_CALL(A) CUDA_SAFE_CALL_ALWAYS(A)
#else
# define CUDA_SAFE_CALL(A) A
#endif
dgsopengl::dgsopengl(QWidget * parent)
    : QOpenGLWidget(parent)
{
    setFocusPolicy(Qt::StrongFocus);
    setFocus(); 
    //三角形测试
    shaderProgram = nullptr;
    VAO = 0;
    VBO = 0;
}
float sigmoid(const float m1)
{
    return 1.0f / (1.0f + exp(-m1));
}

std::function<char* (size_t N)> resizeFunctional(void** ptr, size_t& S) {
    auto lambda = [ptr, &S](size_t N) {
        if (N > S)
        {
            if (*ptr)
                CUDA_SAFE_CALL(cudaFree(*ptr));
            CUDA_SAFE_CALL(cudaMalloc(ptr, 2 * N));
            S = 2 * N;
        }
        return reinterpret_cast<char*>(*ptr);
    };
    return lambda;
}
 
template<int D>
int loadPly(const QString &filename,
    std::vector<Pos>& pos,
    std::vector<SHs<3>>& shs,
    std::vector<float>& opacities,
    std::vector<Scale>& scales,
    std::vector<Rot>& rot,
    QVector3D& minn,
    QVector3D& maxx)
{
    QByteArray ba =filename.toUtf8();
    char* file=ba.data();
    std::ifstream infile(file, std::ios_base::binary);

    if (!infile.good())
      qDebug()<<"Unable to find model's PLY file, attempted"<< filename;

    // "Parse" header (it has to be a specific format anyway)
    std::string buff;
    std::getline(infile, buff);
    std::getline(infile, buff);

    std::string dummy;
    std::getline(infile, buff);
    std::stringstream ss(buff);
    int count;
    //跳过文件头部的三行内容，并读取包含点数（count）的行。count 表示文件中高斯分布（Gaussian Splat）的数量。
    ss >> dummy >> dummy >> count;

    // Output number of Gaussians contained
    std::cout << "Loading " << count << " Gaussian splats" << std::endl;

    while (std::getline(infile, buff))
        if (buff.compare("end_header") == 0)
            break;

    // Read all Gaussians at once (AoS) 读取文件中的高斯点数据（AoS）
    std::vector<RichPoint<D>> points(count);
    infile.read((char*)points.data(), count * sizeof(RichPoint<D>));
    // 分配空间
    pos.resize(count);
    shs.resize(count);
    scales.resize(count);
    rot.resize(count);
    opacities.resize(count);
    minn = QVector3D(FLT_MAX, FLT_MAX, FLT_MAX);
    maxx = QVector3D(-FLT_MAX, -FLT_MAX, -FLT_MAX);
    // 计算边界
    for (int i = 0; i < count; ++i)
    {
        const QVector3D& p = points[i].pos;
        maxx.setX(std::max(maxx.x(), p.x()));
        maxx.setY(std::max(maxx.y(), p.y()));
        maxx.setZ(std::max(maxx.z(), p.z()));

        minn.setX(std::min(minn.x(), p.x()));
        minn.setY(std::min(minn.y(), p.y()));
        minn.setZ(std::min(minn.z(), p.z()));
    }
    qDebug() << "minn:" << minn << "maxx:" << maxx;

    // Morton 排序
    QVector<QPair<uint64_t, int>> mapp(count);
    for (int i = 0; i < count; ++i)
    {
        QVector3D rel = (points[i].pos - minn);
        rel.setX(rel.x() / (maxx.x() - minn.x()));
        rel.setY(rel.y() / (maxx.y() - minn.y()));
        rel.setZ(rel.z() / (maxx.z() - minn.z()));

        QVector3D scaled = rel * float((1 << 21) - 1);
        int xyz[3] = { int(scaled.x()), int(scaled.y()), int(scaled.z()) };

        uint64_t code = 0;
        for (int j = 0; j < 21; ++j)
        {
            code |= ((uint64_t(xyz[0] & (1 << j))) << (2 * j + 0));
            code |= ((uint64_t(xyz[1] & (1 << j))) << (2 * j + 1));
            code |= ((uint64_t(xyz[2] & (1 << j))) << (2 * j + 2));
        }
        mapp[i] = qMakePair(code, i);
    }

    std::sort(mapp.begin(), mapp.end(), [](const QPair<uint64_t, int>& a, const QPair<uint64_t, int>& b) {
        return a.first < b.first;
        });

    int SH_N = (D + 1) * (D + 1);
    for (int k = 0; k < count; ++k)
    {
        int i = mapp[k].second;
        pos[k] = points[i].pos;

        // Normalize quaternion
        float length2 = 0.f;
        for (int j = 0; j < 4; ++j)
            length2 += points[i].rot.rot[j] * points[i].rot.rot[j];
        float length = std::sqrt(length2);
        for (int j = 0; j < 4; ++j)
            rot[k].rot[j] = points[i].rot.rot[j] / length;

        // Exponentiate scale
        for (int j = 0; j < 3; ++j)
            scales[k].scale[j] = std::exp(points[i].scale.scale[j]);

        // Activate alpha
        opacities[k] = sigmoid(points[i].opacity);

        shs[k].shs[0] = points[i].shs.shs[0];
        shs[k].shs[1] = points[i].shs.shs[1];
        shs[k].shs[2] = points[i].shs.shs[2];
        for (int j = 1; j < SH_N; ++j)
        {
            shs[k].shs[j * 3 + 0] = points[i].shs.shs[(j - 1) + 3];
            shs[k].shs[j * 3 + 1] = points[i].shs.shs[(j - 1) + SH_N + 2];
            shs[k].shs[j * 3 + 2] = points[i].shs.shs[(j - 1) + 2 * SH_N + 1];
        }
    }
    qDebug() << "First point position:" << points[0].pos;
    return count;
}


void dgsopengl::initializeGL()
{
    initializeOpenGLFunctions();
    _renderTarget = new RenderTarget(1200, 892);
    _camera = new Camera(1200, 892);
    _monoRdrMode = new MonoRdrMode();
    render_w = _renderTarget->m_W;
    render_h = _renderTarget->m_H;
    _copyRenderer = new BufferCopyRenderer();

    ////三角立体测试
    //glEnable(GL_DEPTH_TEST);
    //glClearColor(0.1f, 0.1f, 0.1f, 1.0f);

    //const char* vshader = R"(
    //    #version 330 core
    //    layout (location = 0) in vec3 aPos;
    //    uniform mat4 model;
    //    uniform mat4 view;
    //    uniform mat4 projection;
    //    void main() {
    //        gl_Position = projection * view * model * vec4(aPos, 1.0);
    //    }
    //)";

    //const char* fshader = R"(
    //    #version 330 core
    //    out vec4 FragColor;
    //    void main() {
    //        FragColor = vec4(0.8, 0.3, 0.2, 1.0);
    //    }
    //)";

    //shaderProgram = new QOpenGLShaderProgram();
    //shaderProgram->addShaderFromSourceCode(QOpenGLShader::Vertex, vshader);
    //shaderProgram->addShaderFromSourceCode(QOpenGLShader::Fragment, fshader);
    //shaderProgram->link();

    //// 一个立体三角形（四面体的 3 个顶点）
    //float vertices[] = {
    //     0.0f,  0.5f,  0.0f,
    //    -0.5f, -0.5f,  0.5f,
    //     0.5f, -0.5f,  0.5f,

    //     0.0f,  0.5f,  0.0f,
    //     0.5f, -0.5f,  0.5f,
    //     0.5f, -0.5f, -0.5f,

    //     0.0f,  0.5f,  0.0f,
    //     0.5f, -0.5f, -0.5f,
    //    -0.5f, -0.5f, -0.5f,

    //     0.0f,  0.5f,  0.0f,
    //    -0.5f, -0.5f, -0.5f,
    //    -0.5f, -0.5f,  0.5f,
    //};

    //glGenVertexArrays(1, &VAO);
    //glGenBuffers(1, &VBO);

    //glBindVertexArray(VAO);
    //glBindBuffer(GL_ARRAY_BUFFER, VBO);
    //glBufferData(GL_ARRAY_BUFFER, sizeof(vertices), vertices, GL_STATIC_DRAW);

    //glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 3 * sizeof(float), (void*)0);
    //glEnableVertexAttribArray(0);
    //glBindBuffer(GL_ARRAY_BUFFER, 0);
    //glBindVertexArray(0);
    //三角立体测试
}

void dgsopengl::resizeGL(int w, int h)
{
    glViewport(0, 0, w, h);
   _renderTarget->resize(w, h);
   render_w = w;  
   render_h = h;  
   // 标记需要更新缓冲区
   _needsBufferUpdate = true;
}

void dgsopengl::paintGL()
{
    // 清屏一次
    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
    glClearColor(0.75f, 0.75f, 0.75f, 0.4f);
    // 渲染 PLY 或其他内容
    if (_plyLoaded) {
        _renderTarget->bind();
        onRenderIBR(*_renderTarget, *_camera);
        _renderTarget->unbind();
        glDisable(GL_BLEND); // 如果不需要混合
    }

    // 绘制立体三角形
  /*  glEnable(GL_DEPTH_TEST);
    shaderProgram->bind();
    glBindVertexArray(VAO);

    QMatrix4x4 model;
    model.setToIdentity();
    QMatrix4x4 view = _camera->getView();
    QMatrix4x4 projection = _camera->getProjection();

    shaderProgram->setUniformValue("model", model);
    shaderProgram->setUniformValue("view", view);
    shaderProgram->setUniformValue("projection", projection);

    glDrawArrays(GL_TRIANGLES, 0, 12);

    glBindVertexArray(0);
    shaderProgram->release();*/

    //三角立体绘制
}

void dgsopengl::loadPlyAndUpload()
{
    qDebug() << "loadPlyandUpload";
    _copyRenderer->flip() = true;
    _copyRenderer->width() = render_w;
    _copyRenderer->height() = render_h;

    if (_fileply.isEmpty()) {
        return;
    }
    // Allocate and fill the GPU data
    // 1. 为位置数据分配 GPU 内存并拷贝数据
    std::vector<Pos> pos;
    std::vector<Rot> rot;
    std::vector<Scale> scale;
    std::vector<float> opacity;
    std::vector<SHs<3>> shs;
    qDebug() << "Loading PLY file:" << _fileply;
    count = loadPly<0>(_fileply, pos, shs, opacity, scale, rot, _scenemin, _scenemax);
    qDebug() << "After call -> _scenemin:" << _scenemin << " _scenemax:" << _scenemax;
    qDebug() << "count==" << count;
    qDebug()<<"pos=="<<pos.data()<<"shs"<<shs.data()<<"opacity"<<opacity.data()<<"scale"<<scale.data()<<"rot"<<rot.data();
    int P = count;
    _boxmin = _scenemin;
    _boxmax = _scenemax;

    CUDA_SAFE_CALL_ALWAYS(cudaMalloc((void**)&pos_cuda, sizeof(Pos) * P));
    CUDA_SAFE_CALL_ALWAYS(cudaMemcpy(pos_cuda, pos.data(), sizeof(Pos) * P, cudaMemcpyHostToDevice));

    // 2. 为旋转数据分配 GPU 内存并拷贝数据
    CUDA_SAFE_CALL_ALWAYS(cudaMalloc((void**)&rot_cuda, sizeof(Rot) * P));
    CUDA_SAFE_CALL_ALWAYS(cudaMemcpy(rot_cuda, rot.data(), sizeof(Rot) * P, cudaMemcpyHostToDevice));

    // 3. 为球谐系数数据分配 GPU 内存并拷贝数据
    CUDA_SAFE_CALL_ALWAYS(cudaMalloc((void**)&shs_cuda, sizeof(SHs<3>) * P));
    CUDA_SAFE_CALL_ALWAYS(cudaMemcpy(shs_cuda, shs.data(), sizeof(SHs<3>) * P, cudaMemcpyHostToDevice));

    // 4. 为不透明度数据分配 GPU 内存并拷贝数据
    CUDA_SAFE_CALL_ALWAYS(cudaMalloc((void**)&opacity_cuda, sizeof(float) * P));
    CUDA_SAFE_CALL_ALWAYS(cudaMemcpy(opacity_cuda, opacity.data(), sizeof(float) * P, cudaMemcpyHostToDevice));

    // 5. 为缩放数据分配 GPU 内存并拷贝数据
    CUDA_SAFE_CALL_ALWAYS(cudaMalloc((void**)&scale_cuda, sizeof(Scale) * P));
    CUDA_SAFE_CALL_ALWAYS(cudaMemcpy(scale_cuda, scale.data(), sizeof(Scale) * P, cudaMemcpyHostToDevice));

    // Create space for view parameters
    CUDA_SAFE_CALL_ALWAYS(cudaMalloc((void**)&view_cuda, sizeof(QMatrix4x4)));
    CUDA_SAFE_CALL_ALWAYS(cudaMalloc((void**)&proj_cuda, sizeof(QMatrix4x4)));
    CUDA_SAFE_CALL_ALWAYS(cudaMalloc((void**)&cam_pos_cuda, 3 * sizeof(float)));
    CUDA_SAFE_CALL_ALWAYS(cudaMalloc((void**)&background_cuda, 3 * sizeof(float)));
    CUDA_SAFE_CALL_ALWAYS(cudaMalloc((void**)&rect_cuda, 2 * P * sizeof(int)));


    float bg[3] = { 0.1, 0.1, 0.1 };
    CUDA_SAFE_CALL_ALWAYS(cudaMemcpy(background_cuda, bg, 3 * sizeof(float), cudaMemcpyHostToDevice));
    /*gData = new GaussianData(P,
        (float*)pos.data(),
        (float*)rot.data(),
        (float*)scale.data(),
        opacity.data(),
        (float*)shs.data());*/
    //_gaussianRenderer = new GaussianSurfaceRenderer();  //---------->GaussianData和GaussianSurfaceRenderer暂时没用
    // Create GL buffer ready for CUDA/GL interop
   
    glCreateBuffers(1, &imageBuffer);
    glNamedBufferStorage(imageBuffer, render_w * render_h * 3 * sizeof(float), nullptr, GL_DYNAMIC_STORAGE_BIT);
    if (cudaPeekAtLastError() != cudaSuccess)
    {
        qDebug() << "A CUDA error occurred in setup:" << cudaGetErrorString(cudaGetLastError()) << ". Please rerun in Debug to find the exact line!";
    }
    cudaError_t err = cudaGraphicsGLRegisterBuffer(&imageBufferCuda, imageBuffer, cudaGraphicsRegisterFlagsWriteDiscard);
    if (err != cudaSuccess) {
        qDebug()<< "Failed to register GL buffer for CUDA: " << cudaGetErrorString(err) << "\n";
    }

    geomBufferFunc = resizeFunctional(&geomPtr, allocdGeom);
    binningBufferFunc = resizeFunctional(&binningPtr, allocdBinning);
    imgBufferFunc = resizeFunctional(&imgPtr, allocdImg);
    _plyLoaded = true;
}

void dgsopengl::onRenderIBR(RenderTarget& dst, const Camera& eye)
{
    //检查是否需要更新缓冲区
    if (_needsBufferUpdate) {
        updateImageBuffer();
        _needsBufferUpdate = false;
    }
   qDebug()<<"render_w==?"<<render_w;
   qDebug()<<"render_h==?"<<render_h;
    //// Convert view and projection to target coordinate system
    auto view_mat = eye.getView();
    auto proj_mat = eye.getViewProjection();
    qDebug() << "view_mat为" << view_mat;
    qDebug()<< "proj_mat为" << proj_mat;
    qDebug() << "眼睛的位置" << eye.position();
    qDebug() << "FOVY" << eye.fovy();
    qDebug() << "eye.aspect()" << eye.aspect();
    view_mat.row(1) *= -1;
    view_mat.row(2) *= -1;
    proj_mat.row(1) *= -1;

    //// Compute additional view parameters
    float tan_fovy = tan(qDegreesToRadians(eye.fovy()) * 0.5f);
    qDebug() << "eye.fovy弧度" << qDegreesToRadians(eye.fovy());
    float tan_fovx = tan_fovy * eye.aspect();

    //// Copy frame-dependent data to GPU
    CUDA_SAFE_CALL(cudaMemcpy(view_cuda, view_mat.data(), sizeof(QMatrix4x4), cudaMemcpyHostToDevice));
    qDebug() << "view_mat.data()" << view_mat.data()<<"proj_mat.data()"<<proj_mat.data();
    CUDA_SAFE_CALL(cudaMemcpy(proj_cuda, proj_mat.data(), sizeof(QMatrix4x4), cudaMemcpyHostToDevice));
    CUDA_SAFE_CALL(cudaMemcpy(cam_pos_cuda, &eye.position(), sizeof(float) * 3, cudaMemcpyHostToDevice));
    float* image_cuda = nullptr;
    size_t bytes;
    // Map OpenGL buffer resource for use with CUDA
    CUDA_SAFE_CALL(cudaGraphicsMapResources(1, &imageBufferCuda));
    CUDA_SAFE_CALL(cudaGraphicsResourceGetMappedPointer((void**)&image_cuda, &bytes, imageBufferCuda));
    qDebug() << "image_cuda=" << image_cuda;
   
    //// Rasterize
    int* rects = _fastCulling ? rect_cuda : nullptr;
    float* boxmin = _cropping ? (float*)&_boxmin : nullptr;
    float* boxmax = _cropping ? (float*)&_boxmax : nullptr;

    CudaRasterizer::Rasterizer::forward(
        geomBufferFunc,
        binningBufferFunc,
        imgBufferFunc,
        count, 0, 16,
        background_cuda,
        render_w, render_h,
        pos_cuda,
        shs_cuda,
        nullptr,
        opacity_cuda,
        scale_cuda,
        _scalingModifier,
        rot_cuda,
        nullptr,
        view_cuda,
        proj_cuda,
        cam_pos_cuda,
        tan_fovx,
        tan_fovy,
        false,
        image_cuda,
        _antialiasing,
        nullptr,
        rects,
        boxmin,
        boxmax
    );
     
    CUDA_SAFE_CALL(cudaGraphicsUnmapResources(1, &imageBufferCuda));
   
    //// Copy image contents to framebuffer
     _copyRenderer->process(imageBuffer, dst, render_w, render_h);  //_resolution.x(), _resolution.y()

    if (cudaPeekAtLastError() != cudaSuccess)
    {
        qDebug() << "A CUDA error occurred during rendering:" << cudaGetErrorString(cudaGetLastError()) << ". Please rerun in Debug to find the exact line!";
    }
}

void dgsopengl::updateImageBuffer()
{
    // 释放旧资源
    if (imageBufferCuda) {
        cudaGraphicsUnregisterResource(imageBufferCuda);
        imageBufferCuda = nullptr;
    }
    if (imageBuffer) {
        glDeleteBuffers(1, &imageBuffer);
        imageBuffer = 0;
    }

    // 创建新缓冲区
    glCreateBuffers(1, &imageBuffer);
    const size_t bufferSize = render_w * render_h * 3 * sizeof(float);
    glNamedBufferStorage(imageBuffer, bufferSize, nullptr, GL_DYNAMIC_STORAGE_BIT);

    // 注册到 CUDA
    cudaError_t err = cudaGraphicsGLRegisterBuffer(&imageBufferCuda, imageBuffer,
        cudaGraphicsRegisterFlagsWriteDiscard);
    if (err != cudaSuccess) {
        qFatal("CUDA 缓冲区注册失败: %s", cudaGetErrorString(err));
    }

    // 更新光栅化缓冲区函数
    imgBufferFunc = resizeFunctional(&imgPtr, allocdImg);
}

dgsopengl::~dgsopengl()
{
    cudaFree(pos_cuda);
    cudaFree(rot_cuda);
    cudaFree(scale_cuda);
    cudaFree(opacity_cuda);
    cudaFree(shs_cuda);

    cudaFree(view_cuda);
    cudaFree(proj_cuda);
    cudaFree(cam_pos_cuda);
    cudaFree(background_cuda);
    cudaFree(rect_cuda);
    delete _renderTarget;
    delete _camera;
    delete _copyRenderer;
    //delete _gaussianRenderer;
    //delete gData;
}

//相机移动
void dgsopengl::keyPressEvent(QKeyEvent* event)  {
    float step = 0.5f; // 每次移动的步长
    switch (event->key()) {
    case Qt::Key_W:
        _camera->moveForward(step);
        break;
    case Qt::Key_S:
        _camera->moveBackward(step);
        break;
    case Qt::Key_A:
        _camera->moveLeft(step);
        break;
    case Qt::Key_D:
        _camera->moveRight(step);
        break;
    default:
        QOpenGLWidget::keyPressEvent(event);
    }
    update(); // 触发重绘
}